package eventBusV4

import (
	"errors"
	"sync"
)

const (
	numWorkers = 4
	taskQueueSize = 1024
	initialSubsCapacity = 16
)

var taskPool = sync.Pool{
	New: func() interface{} {
		return &asyncTask{}
	},
}

type Bus struct {
	subNode     sync.Map
	taskQueue   chan asyncTask
	shutdownSig chan struct{}
	closed      bool
	mu          sync.Mutex
}

func NewBus() Bus {
	bus := Bus{
		taskQueue:   make(chan asyncTask, taskQueueSize),
		shutdownSig: make(chan struct{}),
	}
	for i := 0; i < numWorkers; i++ {
		go bus.worker()
	}
	return bus
}

type asyncTask struct {
	sub Sub
	msg interface{}
}

func (b *Bus) worker() {
	for {
		select {
		case task, ok := <-b.taskQueue:
			if !ok {
				return
			}
			task.sub.receive(task.msg)
		case <-b.shutdownSig:
			return
		}
	}
}

func (b *Bus) Subscribe(topic string, sub Sub) {
	n, _ := b.subNode.LoadOrStore(topic, NewNode())
	node := n.(*node)
	node.rw.Lock()
	node.subs = append(node.subs, sub)
	node.rw.Unlock()
}

func (b *Bus) UnSubscribe(topic string, sub Sub) {
	n, ok := b.subNode.Load(topic)
	if !ok {
		return
	}
	node := n.(*node)
	node.removeSub(sub)
	if node.SubsLen() == 0 {
		b.subNode.Delete(topic)
	}
}

func (b *Bus) Publish(topic string, msg interface{}) error {
	n, ok := b.subNode.Load(topic)
	if !ok {
		return errors.New("topic not exist")
	}
	node := n.(*node)

	node.rw.RLock()
	if len(node.subs) == 0 {
		node.rw.RUnlock()
		return nil
	}
	// 优化切片复制
	subs := make([]Sub, len(node.subs))
	copy(subs, node.subs)
	node.rw.RUnlock()

	// 使用对象池获取任务对象
	task := taskPool.Get().(*asyncTask)
	task.msg = msg
	
	for _, sub := range subs {
		task.sub = sub
		select {
		case b.taskQueue <- *task:
			// 成功发送任务
		default:
			// 队列满时直接发送，避免清空队列
			task.sub.receive(task.msg)
		}
	}
	
	// 归还对象到池
	taskPool.Put(task)
	return nil
}

func (b *Bus) drainTaskQueue() {
	for {
		select {
		case <-b.taskQueue:
		default:
			return
		}
	}
}

func (b *Bus) PubFunc(topic string) func(msg interface{}) {
	return func(msg interface{}) {
		b.Publish(topic, msg)
	}
}

func (b *Bus) SubsLen(topic string) (int, error) {
	n, ok := b.subNode.Load(topic)
	if !ok {
		return 0, errors.New("topic not exist")
	}
	node := n.(*node)
	node.rw.RLock()
	defer node.rw.RUnlock()
	return node.SubsLen(), nil
}

func (b *Bus) Close() {
	b.mu.Lock()
	if b.closed {
		b.mu.Unlock()
		return
	}
	b.closed = true
	b.mu.Unlock()

	close(b.shutdownSig)
	close(b.taskQueue)

	b.subNode.Range(func(key, value interface{}) bool {
		node := value.(*node)
		node.rw.Lock()
		for _, sub := range node.subs {
			close(sub.out)
		}
		node.subs = nil
		node.rw.Unlock()
		return true
	})
}

type node struct {
	subs []Sub
	rw   sync.RWMutex
}

func NewNode() *node {
	return &node{
		subs: make([]Sub, 0, initialSubsCapacity),
	}
}

func (n *node) SubsLen() int {
	return len(n.subs)
}

func (n *node) removeSub(s Sub) {
	n.rw.Lock()
	defer n.rw.Unlock()
	
	idx := n.findSubIdx(s)
	if idx < 0 {
		return
	}
	
	copy(n.subs[idx:], n.subs[idx+1:])
	n.subs = n.subs[:len(n.subs)-1]
}

func (n *node) findSubIdx(s Sub) int {
	for idx, sub := range n.subs {
		if sub == s {
			return idx
		}
	}
	return -1
}

type Sub struct {
	out chan interface{}
}

func NewSub() Sub {
	return Sub{
		out: make(chan interface{}, 1024),
	}
}

func (s *Sub) receive(msg interface{}) {
	s.out <- msg
}

func (s *Sub) Out() chan interface{} {
	return s.out
}